import os
import sys
import traci
import json
import time
import datetime
from collections import defaultdict

# 确保SUMO_HOME环境变量已设置
if 'SUMO_HOME' in os.environ:
    tools = os.path.join(os.environ['SUMO_HOME'], 'tools')
    sys.path.append(tools)
else:
    sys.exit("请设置SUMO_HOME环境变量")

class SUMOSimulator:
    def __init__(self, config_file, junctions_file, gui=True):
        """
        初始化SUMO仿真器
        config_file: SUMO配置文件路径 (.sumocfg)
        junctions_file: 路口数据JSON文件路径
        gui: 是否使用GUI模式
        """
        self.config_file = os.path.abspath(config_file)
        self.junctions_file = os.path.abspath(junctions_file)
        self.gui = gui
        self.simulation_started = False
        self.warmup_done = False
        self.warmup_steps = 300
        self.start_time = None
        self.vehicle_counts = defaultdict(lambda: defaultdict(list))
        self.timestamps = defaultdict(lambda: defaultdict(list))
        
        # 验证文件存在
        if not os.path.exists(self.config_file):
            raise FileNotFoundError(f"SUMO配置文件未找到: {self.config_file}")
        if not os.path.exists(self.junctions_file):
            raise FileNotFoundError(f"路口数据文件未找到: {self.junctions_file}")
        
        # 加载路口配置
        try:
            with open(self.junctions_file, 'r', encoding='utf-8') as f:
                self.junctions_data = json.load(f)
        except Exception as e:
            print(f"Error loading junctions data: {str(e)}")
            self.junctions_data = {}

    def is_connected(self):
        """检查是否连接到SUMO"""
        try:
            return traci.getConnection() is not None
        except:
            return False

    def start_simulation(self):
        """启动仿真并进行预热"""
        try:
            if not self.simulation_started:
                # 如果已经有连接，先关闭
                try:
                    if self.is_connected():
                        traci.close()
                except:
                    pass
                
                # 设置SUMO命令
                sumo_binary = "sumo-gui" if self.gui else "sumo"
                sumo_cmd = [
                    sumo_binary,
                    "-c", self.config_file,
                    "--step-length", "1.0",
                    "--no-warnings", "true",
                    "--start",
                    "--quit-on-end",
                ]
                
                print(f"Starting SUMO with command: {' '.join(sumo_cmd)}")
                
                # 启动SUMO
                traci.start(sumo_cmd)
                self.simulation_started = True
                print("Successfully connected to SUMO")
                
                print("Starting warmup phase...")
                # 预热阶段
                for i in range(self.warmup_steps):
                    if not self.is_connected():
                        raise Exception("SUMO connection lost during warmup")
                    traci.simulationStep()
                    if i % 100 == 0:
                        print(f"Warmup progress: {i}/{self.warmup_steps}")
                
                self.warmup_done = True
                self.start_time = time.time()
                print("Warmup completed. Starting real-time simulation.")
                return True
                
        except Exception as e:
            print(f"Error starting simulation: {str(e)}")
            self.simulation_started = False
            self.warmup_done = False
            try:
                if self.is_connected():
                    traci.close()
            except:
                pass
            return False

    def step(self):
        """执行一个仿真步骤"""
        try:
            if self.simulation_started and self.is_connected():
                traci.simulationStep()
                return True
            return False
        except Exception as e:
            print(f"Error during simulation step: {str(e)}")
            return False
    def get_junction_vehicle_counts(self, junction_id):
        """获取指定路口各方向的交通信息"""
        try:
            if not self.simulation_started or not self.is_connected():
                return None, None
            
            current_time = datetime.datetime.now()
            traffic_data = defaultdict(lambda: {
                'vehicle_count': 0,     # 车辆总数
                'mean_speed': 0,        # 平均速度
                'halting_count': 0,     # 停止的车辆数
                'waiting_time': 0,      # 平均等待时间
            })
            
            junction_info = self.junctions_data.get(junction_id)
            if not junction_info:
                return None, None
            
            for direction, lanes in junction_info["incoming_lanes"].items():
                total_speed = 0
                total_waiting_time = 0
                total_vehicles = 0
                
                for lane in lanes:
                    lane_id = lane['lane_id']
                    try:
                        # 获取车道上的所有车辆ID
                        vehicle_ids = traci.lane.getLastStepVehicleIDs(lane_id)
                        
                        # 基础统计
                        traffic_data[direction]['vehicle_count'] += len(vehicle_ids)
                        traffic_data[direction]['halting_count'] += traci.lane.getLastStepHaltingNumber(lane_id)
                        
                        # 计算平均速度和等待时间
                        for vehicle_id in vehicle_ids:
                            vehicle_speed = traci.vehicle.getSpeed(vehicle_id)
                            waiting_time = traci.vehicle.getAccumulatedWaitingTime(vehicle_id)
                            total_speed += vehicle_speed
                            total_waiting_time += waiting_time
                            total_vehicles += 1
                        
                    except traci.exceptions.TraCIException as e:
                        print(f"Error getting data for lane {lane_id}: {str(e)}")
                        continue
                
                # 计算平均值
                if total_vehicles > 0:
                    traffic_data[direction]['mean_speed'] = total_speed / total_vehicles
                    traffic_data[direction]['waiting_time'] = total_waiting_time / total_vehicles
                
                # 更新历史数据
                self.vehicle_counts[junction_id][direction].append(traffic_data[direction]['vehicle_count'])
                self.timestamps[junction_id][direction].append(current_time)
                
                if len(self.vehicle_counts[junction_id][direction]) > 20:
                    self.vehicle_counts[junction_id][direction].pop(0)
                    self.timestamps[junction_id][direction].pop(0)
            
            return dict(traffic_data), self.get_historical_data(junction_id) # {'南北方向直行与右转': {'vehicle_count': 37, 'mean_speed': 1.0774107930032664, 'halting_count': 30, 'waiting_time': 54.67567567567568}, '南北方向左转': {'vehicle_count': 7, 'mean_speed': 0.5072661575047395, 'halting_count': 6, 'waiting_time': 27.285714285714285}, '东向西通行': {'vehicle_count': 0, 'mean_speed': 0, 'halting_count': 0, 'waiting_time': 0}, '西向东通行': {'vehicle_count': 1, 'mean_speed': 20.99044959471956, 'halting_count': 0, 'waiting_time': 0.0}})
        
            
        except Exception as e:
            print(f"Error getting traffic data: {str(e)}")
            return None, None
    
    def get_historical_data(self, junction_id):
        """获取历史数据"""
        data = []
        try:
            if junction_id not in self.junctions_data:
                return None
            
            for direction in self.vehicle_counts[junction_id]:
                for t, c in zip(self.timestamps[junction_id][direction],
                              self.vehicle_counts[junction_id][direction]):
                    data.append({
                        'timestamp': t,
                        'vehicle_count': c,
                        'direction': direction
                    })
            
            return data
        except Exception as e:
            print(f"Error getting historical data: {str(e)}")
            return None
        
    # 在SUMOSimulator类中添加以下函数

    def get_current_phase(self, junction_id):
        """获取指定路口的当前信号灯相位"""
        try:
            tls_id = junction_id
            current_phase_index = traci.trafficlight.getPhase(tls_id)
            phase_name = self.get_phase_name(tls_id, current_phase_index)
            phase_duration = traci.trafficlight.getPhaseDuration(tls_id)
            remaining_duration = traci.trafficlight.getNextSwitch(tls_id) - traci.simulation.getTime()
            
            return {
                "phase_index": current_phase_index,
                "phase_name": phase_name,
                "total_duration": phase_duration,
                "remaining_duration": remaining_duration
            }
        except Exception as e:
            print(f"Error getting current phase: {str(e)}")
            return None

    def get_phase_name(self, tls_id, phase_index):
        """获取相位的名称描述"""
        try:
            phase_def = traci.trafficlight.getRedYellowGreenState(tls_id)
            # 根据相位状态返回描述性名称
            # 这里需要根据你的信号配置来定制
            phase_names = {
                0: "南北方向直行与右转",
                1: "南北方向左转",
                2: "东向西通行",
                3: "西向东通行"
            }
            return phase_names.get(phase_index, f"Phase {phase_index}")
        except:
            return f"Phase {phase_index}"

    def get_simulation_time(self):
        """
        获取当前仿真时间（秒）
        """
        if self.is_connected():
            try:
                # traci.simulation.getTime() 返回当前仿真步数
                # 将步数转换为秒
                return traci.simulation.getTime()
            except Exception as e:
                print(f"获取仿真时间失败: {str(e)}")
                return 0
        return 0

    def close(self):
        """关闭仿真"""
        try:
            if self.is_connected():
                traci.close()
            self.simulation_started = False
            self.warmup_done = False
        except Exception as e:
            print(f"Error closing simulation: {str(e)}")

def verify_sumo_config(config_file):
    """验证SUMO配置文件"""
    try:
        if not os.path.exists(config_file):
            return False, "配置文件不存在"
        
        # 读取配置文件内容
        with open(config_file, 'r') as f:
            content = f.read()
        
        # 检查必要的配置项
        required_elements = ['net-file', 'route-files']
        for element in required_elements:
            if element not in content:
                return False, f"缺少必要的配置项: {element}"
        
        return True, "配置文件验证通过"
    except Exception as e:
        return False, f"配置文件验证失败: {str(e)}"

def main():
    """测试函数"""
    current_dir = os.getcwd()
    config_file = os.path.join(current_dir, "osm.sumocfg")
    junctions_file = os.path.join(current_dir, "J54_data.json")
    
    is_valid, message = verify_sumo_config(config_file)
    if not is_valid:
        print(f"错误: {message}")
        return
    
    print("正在启动仿真器...")
    simulator = SUMOSimulator(config_file, junctions_file, gui=True)
    
    try:
        print("正在启动仿真...")
        if simulator.start_simulation():
            print("仿真成功启动")
            
            test_junction = 'J54'
            real_name_junction = {'J54': "倪家桥路与领事馆路交叉口"}
            print(f"\n测试路口: {real_name_junction[test_junction]}")
            
            for i in range(5):
                print(f"\n时间步 {i+1}:")
                traffic_data, _ = simulator.get_junction_vehicle_counts(test_junction)
                phase_info = simulator.get_current_phase(test_junction)
                
                print(f"当前信号相位: {phase_info['phase_name']}")
                print(f"剩余时间: {phase_info['remaining_duration']:.1f}秒")
                
                for direction, data in traffic_data.items():
                    print(f"\n方向: {direction}")
                    print(f"  车辆数量: {data['vehicle_count']} 辆")
                    print(f"  平均速度: {data['mean_speed']:.2f} m/s")
                    print(f"  停止车辆数: {data['halting_count']} 辆")
                    print(f"  平均等待时间: {data['waiting_time']:.1f} 秒")
                
                simulator.step()
                time.sleep(1)
                
    except Exception as e:
        print(f"运行时错误: {str(e)}")
    finally:
        print("正在关闭仿真...")
        simulator.close()

if __name__ == "__main__":
    main()